{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/examples/customization/prompts/chat_prompts.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat Prompts Customization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Setup\n",
    "\n",
    "Below, we take the default prompts and customize them to always answer, even if the context is not helpful.\n",
    "\n",
    "We show two ways of setting up the prompts:\n",
    "1. Explicitly define ChatMessage and MessageRole objects.\n",
    "2. Call ChatPromptTemplate.from_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_prompt_str = (\n",
    "    \"Context information is below.\\n\"\n",
    "    \"---------------------\\n\"\n",
    "    \"{context_str}\\n\"\n",
    "    \"---------------------\\n\"\n",
    "    \"Given the context information and not prior knowledge, \"\n",
    "    \"answer the question: {query_str}\\n\"\n",
    ")\n",
    "\n",
    "refine_prompt_str = (\n",
    "    \"We have the opportunity to refine the original answer \"\n",
    "    \"(only if needed) with some more context below.\\n\"\n",
    "    \"------------\\n\"\n",
    "    \"{context_msg}\\n\"\n",
    "    \"------------\\n\"\n",
    "    \"Given the new context, refine the original answer to better \"\n",
    "    \"answer the question: {query_str}. \"\n",
    "    \"If the context isn't useful, output the original answer again.\\n\"\n",
    "    \"Original Answer: {existing_answer}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Explicitly Define `ChatMessage` and `MessageRole` objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage, MessageRole\n",
    "from llama_index.core import ChatPromptTemplate\n",
    "\n",
    "# Text QA Prompt\n",
    "chat_text_qa_msgs = [\n",
    "    ChatMessage(\n",
    "        role=MessageRole.SYSTEM,\n",
    "        content=(\n",
    "            \"Always answer the question, even if the context isn't helpful.\"\n",
    "        ),\n",
    "    ),\n",
    "    ChatMessage(role=MessageRole.USER, content=qa_prompt_str),\n",
    "]\n",
    "text_qa_template = ChatPromptTemplate(chat_text_qa_msgs)\n",
    "\n",
    "# Refine Prompt\n",
    "chat_refine_msgs = [\n",
    "    ChatMessage(\n",
    "        role=MessageRole.SYSTEM,\n",
    "        content=(\n",
    "            \"Always answer the question, even if the context isn't helpful.\"\n",
    "        ),\n",
    "    ),\n",
    "    ChatMessage(role=MessageRole.USER, content=refine_prompt_str),\n",
    "]\n",
    "refine_template = ChatPromptTemplate(chat_refine_msgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. Call `ChatPromptTemplate.from_messages`\n",
    "\n",
    "`from_messages` is syntatic sugar that allows you to define a chat prompt template as a list of tuples, with each tuple corresponding to a chat message (\"role\", \"message\"). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import ChatPromptTemplate\n",
    "\n",
    "# Text QA Prompt\n",
    "chat_text_qa_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        \"Always answer the question, even if the context isn't helpful.\",\n",
    "    ),\n",
    "    (\"user\", qa_prompt_str),\n",
    "]\n",
    "text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)\n",
    "\n",
    "# Refine Prompt\n",
    "chat_refine_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        \"Always answer the question, even if the context isn't helpful.\",\n",
    "    ),\n",
    "    (\"user\", refine_prompt_str),\n",
    "]\n",
    "refine_template = ChatPromptTemplate.from_messages(chat_refine_msgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the Prompts\n",
    "\n",
    "Now, we use the prompts in an index query!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()\n",
    "\n",
    "# Create an index using a chat model, so that we can use the chat prompts!\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Before Adding Templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm unable to provide an answer to that question based on the context information provided.\n"
     ]
    }
   ],
   "source": [
    "print(index.as_query_engine(llm=llm).query(\"Who is Joe Biden?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### After Adding Templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Joe Biden is the current President of the United States, having taken office in January 2021. He previously served as Vice President under President Barack Obama from 2009 to 2017.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    index.as_query_engine(\n",
    "        text_qa_template=text_qa_template,\n",
    "        refine_template=refine_template,\n",
    "        llm=llm,\n",
    "    ).query(\"Who is Joe Biden?\")\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v3",
   "language": "python",
   "name": "llama_index_v3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
